# Internal Includes
from .fgsm import fgsm
from .pgd import pgd




class Attack_builder():
    def __init__(self, attack_name, model, spr, input_size=128):
        self.attack_name = attack_name
        self.model = model
        self.spr = spr
        self.input_size = input_size

    def apply(self, x, y, PL=False, k=5, spr=None, args=None):
        if spr is not None:
            self.spr = spr
        if self.attack_name.upper() == "FGSM":
            return fgsm(x, y, self.model, self.spr, self.input_size, PL)
        if self.attack_name.upper() == "PGD":
            return pgd(x, y, self.model, self.spr, k, self.input_size, PL=PL)
        else:
            raise ValueError("Unknown attack ({})".format(self.attack_name))
